#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright (C) 2023 reinovo, Inc. All Rights Reserved 
#
# @Time    : 2024/4/15 上午11:22
# @Author  : hmm
# @Email   : liuyuhang0531@foxmail.com
# @File    : recognition.py.py
# @Software: PyCharm

import cv2 as cv
import numpy as np
import torch

class Visual_recognition():
    def __init__(self,yolo_path,weight_path,conf=0.5):
        # 初始化模型
        self.model = torch.hub.load(yolo_path, 'custom', path=weight_path, source='local')
        self.model.conf = conf
        # 设置随机颜色区分目标
        self.color = [list(np.random.choice(range(256), size=3)) for i in range(len(self.model.names))]

    def image_detect(self, frame):
        # 预测结果(输入图像格式为rgb)
        results = self.model(frame)
        # 绘制框图，并发布图像
        return self.box_selection(results.pandas().xyxy[0].values, frame)

    def box_selection(self, datas, image):
        # 转换为opencv图像格式bgr
        pub_yolo_data = []
        # 遍历数据
        for data in datas:
            # 发布数据
            yolo_data = dict()
            yolo_data['class_id'] = data[-2]
            yolo_data['class_name'] = data[-1]
            yolo_data['xmin'] = data[0]
            yolo_data['ymin'] = data[1]
            yolo_data['xmax'] = data[2]
            yolo_data['ymax'] = data[3]
            yolo_data['conf'] = data[-3]
            pub_yolo_data.append(yolo_data)
            # print(yolo_data)

            color = self.color[data[-2]]

            # 绘制选择框
            cv.rectangle(image, (int(data[0]), int(data[1] - 3)),
                         (int(data[2]), int(data[3])),
                         (int(color[0]), int(color[1]), int(color[2])),
                         4)
            # 显示类别名
            cv.putText(image, data[-1],
                       (int(data[0]), int(data[1]) - 6),
                       cv.FONT_HERSHEY_SIMPLEX, 0.6, (int(color[0]), int(color[1]), int(color[2])), 1, cv.LINE_AA)
            # 绘制概率
            cv.putText(image, str(int(data[-3] * 100) / 100),
                       (int(data[0]), int(data[3])),
                       cv.FONT_HERSHEY_SIMPLEX, 0.6, (int(color[0]), int(color[1]), int(color[2])), 1, cv.LINE_AA)

            cv.circle(image, (int(data[0] + (data[2] - data[0]) / 2), int(data[1] + (data[3] - data[1]) / 2)), 2,
                      (int(color[0]), int(color[1]), int(color[2])), 2)

        # cv.imshow('image',cv.cvtColor(image, cv.COLOR_RGB2BGR))
        # cv.waitKey(3)
        return image,pub_yolo_data

